// "Copyright [2021] <Copyright Shengkai Lin>"

#include <queue>
#include "network_controller.h"
#include "task_scheduler/task_generator.h"

namespace ai_sim {

Define_Module(NetworkController);

void NetworkController::initialize() {
    clustertype = ClusterType(par("clustertype").intValue());
    cModule * topo_module = getModuleByPath("^.topo_manager");
    topos = omnetpp::check_and_cast<TopoManager *> (topo_module);
    tasgen = omnetpp::check_and_cast<TaskGenerator *> (getModuleByPath("^.taskgenerator"));
    // Initialize random variable generator.
    srand(0);
}





void NetworkController::gen_route(taskboard& tmp_taskboard, RouteAlgorithm* algo) {
////////////////////// Generate and allocate routes/////////////////////////
// Find the shortest path of the comm_pair one by one
// The time complexity of algorithm is pretty HIGH.
    algo->generateRoutes(tmp_taskboard, topos);
    EV << "Finished gen routes" << std::endl;
}





void NetworkController::allocRoutes(const map<int, RouteVec>& routes_map) {
    int tempIndex = 0;
    for (auto tmp_routes : routes_map) {
        if (topos->getControlplaneTopo()->getNodeFor(
            getSimulation()->getModule(tmp_routes.first)) == nullptr) {
            // If the node is not the switch (for example nics)
            continue;
        }
        RouteVecMsg *ra_msg = new RouteVecMsg;
        ra_msg->setKind(ROUTE_ALLOC_MSG);
        ra_msg->setNew_route_table(tmp_routes.second);
        ra_msg->setTargetId(tmp_routes.first);
        // Get the output port to the machine the Gpu in.

        //route_algo = RouteAlgorithm::get(clustertype);
        // If cluster type is ElectricalCluster, controller will send route to adapter instead of oxc or switch directly
        if(clustertype == ClusterType::ElectricalCluster){
            EV<<"send to controller adapter"<<endl;
            ra_msg->setIndex(tempIndex);
            //std::cout<<"yaosi"<<tempIndex<<endl;
            send(ra_msg, "adapter_port");
            EV << "Alloc a route to adapter"<<endl; 
        }
        else{
            int outport = topos->getControlplaneTopo()->getNodeFor(
            getSimulation()->getModule(tmp_routes.first))->
            getinPathfromDest(0)->getRemoteGateId();

            EV << "The remote gate is " << topos->getControlplaneTopo()->getNodeFor(
                getSimulation()->getModule(tmp_routes.first))->
                getinPathfromDest(0)->getRemoteGate()->getFullPath() << std::endl;

            send(ra_msg, outport);
            EV << "Alloc a route to " << getSimulation()->
                getModule(tmp_routes.first)->getFullPath() << std::endl;

        }
        tempIndex += 1;
    }
}

void NetworkController::alloc_taskroutes(taskboard& tmp_taskboard) {
    Enter_Method("alloc_taskroutes(const Task& tmp_task)");


    RouteAlgorithm *route_algo = RouteAlgorithm::get(clustertype);

    gen_route(tmp_taskboard, route_algo);

    // Allocate routes for all stages
    for (auto &temp : tmp_taskboard.stage_info) {
        allocRoutes(temp.routes_on_switch);
    }
    delete route_algo;
}





void NetworkController::flow_ctrl_msg_send(int gpuid, const vector<Flow>& tm_the_flows, Task& the_task)
{
        Enter_Method("flow_ctrl_msg_send");
        TaskFlowVecMsg *tf_msg = new TaskFlowVecMsg;
        for (auto& flow :tm_the_flows) {
            if (flow.outportId > 3670018 || flow.outportId < 0) {
            EV_ERROR << "outport" << flow.outportId << endl;
            EV_ERROR << "Src, Dest" << flow.srcIdx << " , " << flow.destIdx << endl;
            }
        }
        tf_msg->setNow_task(the_task);
        tf_msg->setThe_flows(tm_the_flows);
        tf_msg->setTargetId(gpuid);
        tf_msg->setKind(TASK_ALLOC_MSG);

       cModule *machine_mod = topos->getWorkersTopo()->
            getNode(gpuid)->getModule()->getParentModule();
        topoNode* machine_node =
            topos->getControlplaneTopo()->getNodeFor(machine_mod);
        assert(machine_node->getNumPaths() == 0);
        int outport = machine_node->getinPathfromDest(0)->getRemoteGateId();

        EV << "I am " << machine_mod->getFullPath() << std::endl;
        EV << "My local gate is " << topos->getControlplaneTopo()->
            getNodeFor(machine_mod)-> getinPathfromDest(0)->
            getLocalGate()->getFullPath() << std::endl;
        EV << "The remote gate is " << topos->getControlplaneTopo()->
            getNodeFor(machine_mod)-> getinPathfromDest(0)->
            getRemoteGate()->getFullPath() << std::endl;

        send(tf_msg, outport);
        EV << "Alloc a task to CPU" << gpuid << std::endl;

}


void NetworkController::finish() {
}


void NetworkController::handleMessage(omnetpp::cMessage *msg) {
    if (msg->getKind() == FINISHED_MSG) {
        TaskMsg *ta_msg = dynamic_cast<TaskMsg *>(msg);
	    tasgen->release_task(ta_msg);
        delete ta_msg; //delete ta_msg;
    } else if (msg->getKind() == CHANGE_ROUTE_MSG) {
        // Special serve the HD communication pattern
        FlowMsg* f_msg = dynamic_cast<FlowMsg*>(msg);
        allocRoutes(f_msg->getThe_flow().dependRoute);
        delete f_msg;
    } else if (msg->getKind() == SYNC_MSG) {
        error("DONT SUPPORT SYNC MSG YET");
        // FlowMsg *sy_msg = dynamic_cast<FlowMsg *>(msg);
        // send(msg, "machine_port$o", sy_msg->getDestIdx()%numGpusperMachine);
    } else {
        delete msg;
        error("ERROR MSG ARRIVED ON CONTROLLER");
    }
}

}  // namespace ai_sim
